package com.bruce.ai.alibaba.agentcard.discovery.auth;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.time.Instant;
import java.util.Base64;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * API密钥管理器
 * 负责API密钥的生成、验证和管理
 * 
 * 特性：
 * - API密钥生成和验证
 * - 密钥轮换支持
 * - 使用统计
 * - 安全哈希存储
 */
@Component
public class ApiKeyManager {
    
    private static final Logger logger = LoggerFactory.getLogger(ApiKeyManager.class);
    private static final String HASH_ALGORITHM = "SHA-256";
    private static final int API_KEY_LENGTH = 32;
    
    private final ConcurrentMap<String, ApiKeyInfo> apiKeys;
    private final SecureRandom secureRandom;
    
    public ApiKeyManager() {
        this.apiKeys = new ConcurrentHashMap<>();
        this.secureRandom = new SecureRandom();
    }
    
    /**
     * 生成新的API密钥
     * @param clientId 客户端ID
     * @param description 密钥描述
     * @param expiresInSeconds 过期时间（秒），0表示永不过期
     * @return API密钥
     */
    public String generateApiKey(String clientId, String description, long expiresInSeconds) {
        byte[] keyBytes = new byte[API_KEY_LENGTH];
        secureRandom.nextBytes(keyBytes);
        String apiKey = Base64.getUrlEncoder().withoutPadding().encodeToString(keyBytes);
        
        long expiresAt = expiresInSeconds > 0 ? Instant.now().getEpochSecond() + expiresInSeconds : 0;
        
        ApiKeyInfo keyInfo = new ApiKeyInfo(
                clientId,
                description,
                hashApiKey(apiKey),
                Instant.now().getEpochSecond(),
                expiresAt
        );
        
        apiKeys.put(apiKey, keyInfo);
        
        logger.info("Generated new API key for client: {} ({})", clientId, description);
        return apiKey;
    }
    
    /**
     * 验证API密钥
     * @param apiKey API密钥
     * @return 验证结果
     */
    public ApiKeyValidationResult validateApiKey(String apiKey) {
        if (apiKey == null || apiKey.trim().isEmpty()) {
            return ApiKeyValidationResult.invalid("API key is null or empty");
        }
        
        ApiKeyInfo keyInfo = apiKeys.get(apiKey);
        if (keyInfo == null) {
            return ApiKeyValidationResult.invalid("Invalid API key");
        }
        
        // 检查是否过期
        if (keyInfo.isExpired()) {
            return ApiKeyValidationResult.expired("API key has expired");
        }
        
        // 更新使用统计
        keyInfo.incrementUsage();
        
        return ApiKeyValidationResult.valid(keyInfo);
    }
    
    /**
     * 撤销API密钥
     * @param apiKey API密钥
     * @return 是否成功撤销
     */
    public boolean revokeApiKey(String apiKey) {
        ApiKeyInfo keyInfo = apiKeys.remove(apiKey);
        if (keyInfo != null) {
            logger.info("Revoked API key for client: {} ({})", keyInfo.getClientId(), keyInfo.getDescription());
            return true;
        }
        return false;
    }
    
    /**
     * 获取客户端的所有API密钥信息
     * @param clientId 客户端ID
     * @return API密钥信息列表
     */
    public java.util.List<ApiKeyInfo> getApiKeysForClient(String clientId) {
        return apiKeys.values().stream()
                .filter(keyInfo -> clientId.equals(keyInfo.getClientId()))
                .collect(java.util.stream.Collectors.toList());
    }
    
    /**
     * 清理过期的API密钥
     * @return 清理的密钥数量
     */
    public int cleanupExpiredKeys() {
        int removedCount = 0;
        for (String key : apiKeys.keySet()) {
            ApiKeyInfo keyInfo = apiKeys.get(key);
            if (keyInfo != null && keyInfo.isExpired()) {
                apiKeys.remove(key);
                removedCount++;
            }
        }
        if (removedCount > 0) {
            logger.info("Cleaned up {} expired API keys", removedCount);
        }
        return removedCount;
    }
    
    /**
     * 获取API密钥统计信息
     * @return 统计信息
     */
    public ApiKeyStats getStats() {
        int totalKeys = apiKeys.size();
        int expiredKeys = (int) apiKeys.values().stream().filter(ApiKeyInfo::isExpired).count();
        long totalUsage = apiKeys.values().stream().mapToLong(ApiKeyInfo::getUsageCount).sum();
        
        return new ApiKeyStats(totalKeys, expiredKeys, totalUsage);
    }
    
    /**
     * 轮换API密钥
     * @param oldApiKey 旧的API密钥
     * @param expiresInSeconds 新密钥过期时间（秒）
     * @return 新的API密钥
     */
    public String rotateApiKey(String oldApiKey, long expiresInSeconds) {
        ApiKeyInfo oldKeyInfo = apiKeys.get(oldApiKey);
        if (oldKeyInfo == null) {
            throw new IllegalArgumentException("Old API key not found");
        }
        
        // 生成新密钥
        String newApiKey = generateApiKey(
                oldKeyInfo.getClientId(),
                oldKeyInfo.getDescription() + " (rotated)",
                expiresInSeconds
        );
        
        // 撤销旧密钥
        revokeApiKey(oldApiKey);
        
        logger.info("Rotated API key for client: {}", oldKeyInfo.getClientId());
        return newApiKey;
    }
    
    /**
     * 对API密钥进行哈希处理
     * @param apiKey API密钥
     * @return 哈希值
     */
    private String hashApiKey(String apiKey) {
        try {
            MessageDigest digest = MessageDigest.getInstance(HASH_ALGORITHM);
            byte[] hash = digest.digest(apiKey.getBytes());
            return Base64.getEncoder().encodeToString(hash);
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException("Hash algorithm not available", e);
        }
    }
    
    /**
     * API密钥信息类
     */
    public static class ApiKeyInfo {
        private final String clientId;
        private final String description;
        private final String hashedKey;
        private final long createdAt;
        private final long expiresAt;
        private volatile long usageCount;
        private volatile long lastUsedAt;
        
        public ApiKeyInfo(String clientId, String description, String hashedKey, long createdAt, long expiresAt) {
            this.clientId = clientId;
            this.description = description;
            this.hashedKey = hashedKey;
            this.createdAt = createdAt;
            this.expiresAt = expiresAt;
            this.usageCount = 0;
            this.lastUsedAt = 0;
        }
        
        public String getClientId() {
            return clientId;
        }
        
        public String getDescription() {
            return description;
        }
        
        public String getHashedKey() {
            return hashedKey;
        }
        
        public long getCreatedAt() {
            return createdAt;
        }
        
        public long getExpiresAt() {
            return expiresAt;
        }
        
        public long getUsageCount() {
            return usageCount;
        }
        
        public long getLastUsedAt() {
            return lastUsedAt;
        }
        
        public boolean isExpired() {
            return expiresAt > 0 && Instant.now().getEpochSecond() > expiresAt;
        }
        
        public void incrementUsage() {
            this.usageCount++;
            this.lastUsedAt = Instant.now().getEpochSecond();
        }
    }
    
    /**
     * API密钥验证结果
     */
    public static class ApiKeyValidationResult {
        private final boolean valid;
        private final boolean expired;
        private final String errorMessage;
        private final ApiKeyInfo keyInfo;
        
        private ApiKeyValidationResult(boolean valid, boolean expired, String errorMessage, ApiKeyInfo keyInfo) {
            this.valid = valid;
            this.expired = expired;
            this.errorMessage = errorMessage;
            this.keyInfo = keyInfo;
        }
        
        public static ApiKeyValidationResult valid(ApiKeyInfo keyInfo) {
            return new ApiKeyValidationResult(true, false, null, keyInfo);
        }
        
        public static ApiKeyValidationResult invalid(String errorMessage) {
            return new ApiKeyValidationResult(false, false, errorMessage, null);
        }
        
        public static ApiKeyValidationResult expired(String errorMessage) {
            return new ApiKeyValidationResult(false, true, errorMessage, null);
        }
        
        public boolean isValid() {
            return valid;
        }
        
        public boolean isExpired() {
            return expired;
        }
        
        public String getErrorMessage() {
            return errorMessage;
        }
        
        public ApiKeyInfo getKeyInfo() {
            return keyInfo;
        }
    }
    
    /**
     * API密钥统计信息
     */
    public static class ApiKeyStats {
        private final int totalKeys;
        private final int expiredKeys;
        private final long totalUsage;
        
        public ApiKeyStats(int totalKeys, int expiredKeys, long totalUsage) {
            this.totalKeys = totalKeys;
            this.expiredKeys = expiredKeys;
            this.totalUsage = totalUsage;
        }
        
        public int getTotalKeys() {
            return totalKeys;
        }
        
        public int getExpiredKeys() {
            return expiredKeys;
        }
        
        public int getActiveKeys() {
            return totalKeys - expiredKeys;
        }
        
        public long getTotalUsage() {
            return totalUsage;
        }
    }
}